from langchain.tools import BaseTool
from fastmcp import Client
import asyncio
from functools import partial
from typing import Optional, Type
from pydantic import BaseModel, Field
import json

class MCPToolInput(BaseModel):
    """Input schema for MCP Tool."""
    params: dict = Field(description="Parameters to send to the MCP tool")

class MCPTool(BaseTool):
    name: str = Field(default="")
    description: str = Field(default="")
    args_schema: Type[BaseModel] = MCPToolInput
    tool_name: str = Field(default="")
    server_url: str = Field(default="http://localhost:8000/mcp")
    client: Optional[Client] = Field(default=None)

    def __init__(self, tool_name: str, server_url: str = "http://localhost:8000/mcp", **kwargs):

        # 设置BaseTool需要的属性
        name = f"mcp_{tool_name}"
        description = f"MCP tool for {tool_name}. Use this tool to search or query information using {tool_name}."
        
        # 调用父类初始化
        super().__init__(
            name=name,
            description=description,
            args_schema=MCPToolInput,
            **kwargs
        )
        
        # 设置自定义属性
        self.tool_name = tool_name
        self.server_url = server_url
        self.client = Client(server_url)
        
    def _process_result(self, result) -> str:
        """Process and deserialize the result from MCP tool."""
        if result is None:
            return "No result returned"
        
        # 处理 CallToolResult 类型
        if hasattr(result, 'content'):
            # 提取 content 字段
            content = result.content
            if isinstance(content, list) and len(content) > 0:
                # 如果 content 是列表，取第一个元素
                first_content = content[0]
                if hasattr(first_content, 'text'):
                    return first_content.text
                elif hasattr(first_content, 'data'):
                    return str(first_content.data)
                else:
                    return str(first_content)
            elif hasattr(content, 'text'):
                return content.text
            elif hasattr(content, 'data'):
                return str(content.data)
            else:
                return str(content)
        
        # 如果有 data 属性
        if hasattr(result, 'data'):
            return str(result.data)
        
        # 如果有 text 属性
        if hasattr(result, 'text'):
            return result.text
        
        # 如果结果是字符串且看起来像JSON，尝试解析
        if isinstance(result, str):
            try:
                # 尝试解析JSON并提取实际内容
                parsed = json.loads(result)
                if isinstance(parsed, dict) and 'data' in parsed:
                    return parsed['data']
                elif isinstance(parsed, dict) and 'content' in parsed:
                    return parsed['content']
                else:
                    return str(parsed)
            except json.JSONDecodeError:
                # 如果不是JSON，直接返回原字符串
                return result
        
        # 如果结果是字典，提取有用信息
        if isinstance(result, dict):
            if 'data' in result:
                return str(result['data'])
            elif 'content' in result:
                return str(result['content'])
            else:
                return str(result)
        
        return str(result)


    
    def _run(self, params: dict, run_manager: Optional[any] = None) -> str:
        """Execute the tool with the given parameters."""
        async def _async_call():
            async with self.client:
                result = await self.client.call_tool(self.tool_name, params)
                return result
        
        try:
            result = asyncio.run(_async_call())
            return self._process_result(result)
        except Exception as e:
            return f"Error calling MCP tool: {str(e)}"
    
    async def _arun(self, params: dict, run_manager: Optional[any] = None) -> str:
        """Async version of _run."""
        try:
            async with self.client:
                result = await self.client.call_tool(self.tool_name, params)
                return self._process_result(result)
        except Exception as e:
            return f"Error calling MCP tool: {str(e)}"


bing_tool = MCPTool("bing_search")
gl_forum_tool = MCPTool("gl_forum_search")

if __name__ == "__main__":
    # 正确的调用方式是使用invoke方法，传入参数字典
    try:
        testStr = "ganla\n"
        print(testStr)
        print("1**\n")

        print("测试bing_tool:")
        result1 = bing_tool.invoke({"params": {"query": "你好"}})
        # result1 = result1.replace("\\n", "\n")
        print("bing_tool结果: ")
        print(result1)

        print("==================================================")
        print("\n测试gl_forum_tool:")
        result2 = gl_forum_tool.invoke({"params": {"topic": "你好"}})
        result2 = result2.replace("\\n", "\n")

        print("gl_forum_tool结果: ")
        print(result2)
    except Exception as e:
        print(f"调用出错: {e}")
        print("这可能是因为MCP服务器未运行或连接失败")